from api.controllers.utils.dmdb_utils import get_database_connection, create_sql_chain, execute_sql_query, get_sql_database, DMDatabase
from langchain_openai import ChatOpenAI
import os
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

os.environ["OPENAI_BASE_URL"] = os.getenv("OPENAI_BASE_URL")
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")

def test_dm_database():
    """测试达梦数据库连接和查询"""
    print("测试达梦数据库连接...")
    
    # 测试原始连接
    conn = get_database_connection()
    if conn:
        print("原始数据库连接成功")
        conn.close()
    else:
        print("原始数据库连接失败")
        return
    
    # 测试SQLDatabase对象
    print("\n测试SQLDatabase对象...")
    db = get_sql_database()
    if db and isinstance(db, DMDatabase):
        print("达梦数据库适配器创建成功")
        print(f"数据库方言: {db.dialect}")
        try:
            tables = db.get_usable_table_names()
            print(f"可用表: {tables}")
        except Exception as e:
            print(f"获取表名时出错: {e}")
    else:
        print("SQLDatabase对象创建失败")
        return
    
    # 测试模型
    print("\n测试模型创建...")
    try:
        chat_model = ChatOpenAI(
            model=os.getenv("OPENAI_MODEL", "deepseek-chat"),
            temperature=0
        )
        print("模型创建成功")
    except Exception as e:
        print(f"模型创建失败: {e}")
        return
    
    # 测试SQL查询链
    print("\n测试SQL查询链...")
    chain = create_sql_chain(chat_model)
    if chain:
        print("SQL查询链创建成功")
    else:
        print("SQL查询链创建失败")
        print("注意：create_sql_query_chain 可能不完全支持达梦数据库")
        print("您可以考虑使用自定义实现来处理自然语言到SQL的转换")
        return

if __name__ == "__main__":
    test_dm_database()